package org.apache.sad.core.framework.encryption.impl;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.security.PrivateKey;
import java.security.PublicKey;

import javax.crypto.Cipher;

import org.apache.commons.codec.binary.Base64;
import org.springframework.core.io.DefaultResourceLoader;
import org.springframework.core.io.Resource;

import org.apache.sad.core.framework.encryption.Encryptor;

public class RsaEncryptor implements Encryptor {

	private final static String ENCRYPTION_PREFIX = "{RSA}";
	private final static String CIPHER_ALG = "RSA/ECB/PKCS1Padding";
	private final static String DATA_ENCODING_SET = "UTF-8";
	
	private PrivateKey privateKey;
	private PublicKey publicKey;
	
	public void setPrivateKeyString(String privateKeyString) {
		loadPrivateKey(privateKeyString);
	}
	
	public void setPrivateKeyLocation(String privateKeyLocation) throws IOException, ClassNotFoundException {
		this.privateKey = (PrivateKey)loadKey(new DefaultResourceLoader().getResource(privateKeyLocation));
	}

	public void setPublicKeyLocation(String publicKeyLocation) throws IOException, ClassNotFoundException {
		this.publicKey = (PublicKey)loadKey(new DefaultResourceLoader().getResource(publicKeyLocation));
	}

	private Object loadKey(Resource keyResource) throws IOException, ClassNotFoundException {
		ObjectInputStream ois = new ObjectInputStream(keyResource.getInputStream());
		return ois.readObject();
	}

	public String decrypt(String encryptedData) {
		if ( !isEncrypted(encryptedData) ) {
			return encryptedData;
		}
		encryptedData = encryptedData.substring(ENCRYPTION_PREFIX.length());

		byte[] decryptedBytes;
		try {
			Cipher cipher = Cipher.getInstance(CIPHER_ALG);
			cipher.init(Cipher.DECRYPT_MODE, privateKey);
			byte[] decodedBytes = Base64.decodeBase64(encryptedData.getBytes());
			decryptedBytes = cipher.doFinal(decodedBytes);
			return new String(decryptedBytes, DATA_ENCODING_SET);
		} catch (Exception e) {
			throw new RuntimeException("Unable to decrypt the data using RSA algorithm.", e);
		}
	}

	public String encrypt(String clearData) {
		if ( clearData == null || clearData.length() == 0 || isEncrypted(clearData) ) {
			return clearData;
		}		
		try {
			Cipher cipher = Cipher.getInstance(CIPHER_ALG);
			cipher.init(Cipher.ENCRYPT_MODE, publicKey);		
			byte[] encryptedBytes = cipher.doFinal(clearData.getBytes("UTF-8"));
			return ENCRYPTION_PREFIX + new String(Base64.encodeBase64(encryptedBytes));
		} catch (Exception e) {
			throw new RuntimeException("Unable to encrypt data using RSA algorithm.", e);
		}
	}

	public boolean isEncrypted(String data) {
		if ( data != null && data.startsWith(ENCRYPTION_PREFIX)) {
			return true;
		}
		return false;
	}
	
	private void loadPrivateKey(String privateKeyString) {
		try {
			byte[] privateKeyBytes = Base64.decodeBase64(privateKeyString.getBytes());
			ByteArrayInputStream bais = new ByteArrayInputStream(privateKeyBytes);
			ObjectInputStream ois = new ObjectInputStream(bais);
			this.privateKey = (PrivateKey)ois.readObject();		
		} catch (Exception e) {
			throw new RuntimeException("Unable to load PrivateKey.", e);
		}
	}

}
